from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

# tagging_prompt = ChatPromptTemplate.from_template(
#     """
# Extract the desired information from the following passage.
#
# Only extract the properties mentioned in the 'Classification' function.
#
# Passage:
# {input}
# """
# )


# class Classification(BaseModel):
#     sentiment: str = Field(description="The sentiment of the text")
#     aggressiveness: int = Field(
#         description="How aggressive the text is on a scale from 1 to 10"
#     )
#     language: str = Field(description="The language the text is written in")
#
#
# # LLM
# llm = ChatOpenAI(
#     api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
#     model='deepseek-ai/DeepSeek-V3',
#     base_url="https://www.henapi.top/v1",
#     temperature=0
# ).with_structured_output(
#     Classification
# )
#
# tagging_chain = tagging_prompt | llm
#
# inp = "Estoy increiblemente contento de haberte conocido! Creo que seremos muy buenos amigos!"
# resp=tagging_chain.invoke({"input": inp})
# print(resp.dict())



class Classification(BaseModel):
    sentiment: str = Field(..., enum=["happy", "neutral", "sad"])
    aggressiveness: int = Field(
        ...,
        description="describes how aggressive the statement is, the higher the number the more aggressive",
        enum=[1, 2, 3, 4, 5],
    )
    language: str = Field(
        ..., enum=["spanish", "english", "french", "german", "italian"]
    )

tagging_prompt = ChatPromptTemplate.from_template(
    """
Extract the desired information from the following passage.

Only extract the properties mentioned in the 'Classification' function.

Passage:
{input}
"""
)

# LLM
llm = ChatOpenAI(
    api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
    model='deepseek-ai/DeepSeek-V3',
    base_url="https://www.henapi.top/v1",
    temperature=0
).with_structured_output(
    Classification
)

chain = tagging_prompt | llm


inp = "Estoy increiblemente contento de haberte conocido! Creo que seremos muy buenos amigos!"
resp=chain.invoke({"input": inp})
print(resp.dict())

inp = "Estoy muy enojado con vos! Te voy a dar tu merecido!"
resp=chain.invoke({"input": inp})
print(resp.dict())

inp = "Weather is ok here, I can go outside without much more than a coat"
resp=chain.invoke({"input": inp})
print(resp.dict())